package com.ruoyi.system.chat.chatqwen;

import cn.hutool.core.util.ObjectUtil;
import cn.hutool.core.util.StrUtil;
import cn.hutool.http.Header;
import cn.hutool.http.HttpRequest;
import cn.hutool.http.HttpResponse;
import cn.hutool.json.JSONUtil;
import com.baomidou.mybatisplus.core.toolkit.StringUtils;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.knuddels.jtokkit.Encodings;
import com.knuddels.jtokkit.api.Encoding;
import com.knuddels.jtokkit.api.EncodingRegistry;
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.github.pagehelper.util.StringUtil;
import com.knuddels.jtokkit.api.EncodingType;
import com.ruoyi.common.chat.exception.*;
import com.ruoyi.common.entity.AiModelEntity;
import com.ruoyi.common.entity.AiPlatformEntity;
import com.ruoyi.system.chat.chatqwen.entity.ChatRequestEntity;
import com.ruoyi.system.chat.entity.AiResponseEntity;
import com.ruoyi.system.chat.entity.ApiKeyBalanceEntity;
import com.ruoyi.system.chat.service.ChatService;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.io.OutputStream;
import java.math.BigDecimal;
import java.math.RoundingMode;
import java.nio.charset.StandardCharsets;
import java.time.LocalDate;
import java.time.format.DateTimeFormatter;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.stream.Collectors;

@Service
@Slf4j
public class ChatQwenService implements ChatService {

    @Autowired
    private ObjectMapper objectMapper;
    private static final int PREFIX_LENGTH = "data: ".length();

    private void fillTotal(String apiKey, ApiKeyBalanceEntity balanceEntity, String baseUrl)
            throws JsonProcessingException {
        if (baseUrl == null || baseUrl.trim().isEmpty()) {
            baseUrl = AiPlatformEntity.BASE_URL_OPEN_AI;
        }
        HttpResponse response = HttpRequest.get(baseUrl + "/dashboard/billing/subscription")
                .header(Header.CONTENT_TYPE, "application/json")
                .header(Header.AUTHORIZATION, "Bearer " + apiKey)
                .timeout(15 * 1000)
                .executeAsync();

        int status = response.getStatus();
        String body = response.body();
        balanceEntity.setSubscriptions(body);
        Map<String, Object> map = objectMapper.readValue(body, Map.class);

        log.info("subscription map = {}", map);
        Object errorObj = map.get("error");
        if (errorObj != null) {
            if (errorObj instanceof Map) {
                Map<String, Object> errorMap = (Map<String, Object>) errorObj;
                ApiKeyBalanceEntity.Error error = new ApiKeyBalanceEntity.Error();
                error.setStatus(status);
                error.setCode((String) errorMap.get("code"));
                error.setType((String) errorMap.get("type"));
                balanceEntity.setError(error);
            } else {
                ApiKeyBalanceEntity.Error error = new ApiKeyBalanceEntity.Error();
                error.setStatus(status);
                error.setCode("unknown");
                error.setType("unknown");
                balanceEntity.setError(error);
            }
            return;
        }
        BigDecimal hard_limit_usd = new BigDecimal(String.valueOf(map.get("hard_limit_usd")));
        hard_limit_usd = hard_limit_usd.multiply(new BigDecimal(100));
        hard_limit_usd = hard_limit_usd.divide(new BigDecimal(100), 2, RoundingMode.HALF_UP);
        balanceEntity.setHardLimitUsd(hard_limit_usd);
        balanceEntity.setObject("billing_subscription_and_usage");
        balanceEntity.setHasPaymentMethod((boolean) map.get("has_payment_method"));

    }

    private void fillUsed(String apiKey, ApiKeyBalanceEntity balanceEntity, String baseUrl)
            throws JsonProcessingException {
        if (baseUrl == null || baseUrl.trim().isEmpty()) {
            baseUrl = AiPlatformEntity.BASE_URL_OPEN_AI;
        }
        LocalDate currentDate = LocalDate.now();
        LocalDate firstDayOfMonth = currentDate.withDayOfMonth(1);
        LocalDate lastDayOfMonth = currentDate.withDayOfMonth(currentDate.lengthOfMonth());
        DateTimeFormatter formatter = DateTimeFormatter.ofPattern("yyyy-MM-dd");
        String firstDayOfMonthStr = firstDayOfMonth.format(formatter);
        String lastDayOfMonthStr = lastDayOfMonth.format(formatter);

        HttpResponse response = HttpRequest.get(baseUrl +
                        String.format("/dashboard/billing/usage?start_date=%s&end_date=%s",
                                firstDayOfMonthStr, lastDayOfMonthStr))
                .header(Header.CONTENT_TYPE, "application/json")
                .header(Header.AUTHORIZATION, "Bearer " + apiKey)
                .timeout(15 * 1000)
                .executeAsync();

        String body = response.body();
        balanceEntity.setUsage(body);
        Map<String, Object> map = objectMapper.readValue(body, Map.class);

        log.info("usage map = {}", map);
        BigDecimal total_usage = new BigDecimal(String.valueOf(map.get("total_usage")));
        total_usage = total_usage.divide(new BigDecimal(100), 2, RoundingMode.HALF_UP);
        balanceEntity.setTotalUsage(total_usage);
    }

    @Override
    public AiResponseEntity chat(ChatRequestEntity chatRequestEntity,
                                 HttpServletRequest request,
                                 HttpServletResponse response) throws IOException, ChatException {
        log.info("chat: {}", chatRequestEntity);

        AiModelEntity modelEntity = chatRequestEntity.getModelEntity();
        String key = chatRequestEntity.getApiKeyEntity().getKey();
//        String baseUrl = getBaseUrl(modelEntity);
        String baseUrl = modelEntity.getPath();
        String requestBody = getRequestBody(chatRequestEntity, modelEntity);
        log.info("requestBody: {}", requestBody);
//        int inputTokens = countMessageTokens(null, chatRequestEntity.getMessages());

//        log.info("requestBody: {}, inputTokens = {}", requestBody, inputTokens);

        AiResponseEntity responseEntity = new AiResponseEntity();
//        responseEntity.setInputTokens(inputTokens);
        // 调用接口
        HttpResponse result;
//        UserEntity currentUser = AuthUtils.getCurrentUser();

        OutputStream outputStream = null;
        try {
            result = HttpRequest.post(baseUrl)
                    .header(Header.CONTENT_TYPE, "application/json")
                    .header(Header.AUTHORIZATION, "Bearer " + key)
                    .timeout( 300 * 1000)
                    .body(requestBody)
                    .executeAsync();

            String contentType = result.header(Header.CONTENT_TYPE);
            log.info("contentType[{}]: {}", Header.CONTENT_TYPE, contentType);
            InputStreamReader reader = new InputStreamReader(result.bodyStream());
            outputStream = response.getOutputStream();
            if (contentType != null && contentType.contains("application/json")) {
                // 流式请求却返回json，就认为是错误，不计额度
                log.info("handleWithJson");
                response.setContentType(contentType);
                handleWithJson(reader, outputStream, responseEntity);
            } else if (contentType != null && contentType.contains("text/plain")) {
                log.info("handleWithTextPlain");
                handleWithTextPlain(reader, outputStream, responseEntity);

                int outputTokens = getTokens(responseEntity.getContent());
                outputTokens += 3; // every reply is primed with <|start|>assistant<|message|>
                responseEntity.setOutputTokens(outputTokens);

            } else {
                log.info("handleWithStream");
                response.setContentType("text/event-stream");
                handleWithStream(reader, outputStream, responseEntity);

                int outputTokens = getTokens(responseEntity.getContent());

                outputTokens += 3; // every reply is primed with <|start|>assistant<|message|>

                responseEntity.setOutputTokens(outputTokens);

            }
        } catch (Exception e) {
            log.error("request remote api error: {}", e.getMessage(), e);
            throw new ChatException("request remote api  error: " + e.getMessage(), e);
        } finally {
            try {
                if (outputStream != null) {
                    outputStream.close();
                }
            } catch (Exception e) {
                log.error(e.getMessage(), e);
            }
        }
        log.info("finished");
        return responseEntity;
    }


    private void handleWithJson(InputStreamReader reader, OutputStream outputStream, AiResponseEntity responseEntity) throws IOException {
        int length;
        char[] buffer = new char[1024];
        StringBuilder builder = new StringBuilder();
        while ((length = reader.read(buffer, 0, buffer.length)) != -1) {
            String resp = new String(buffer, 0, length);
            log.info("resp: {}", resp);
            builder.append(resp);
            outputStream.write(resp.getBytes(StandardCharsets.UTF_8));
            outputStream.flush();
        }
        String message = builder.toString();
        try {
            Map<String, Object> map = objectMapper.readValue(message, new TypeReference<Map<String, Object>>() {
            });
            if (map != null && map.get("choices") != null) {
                List<Object> choices = (List<Object>) map.get("choices");
                if (choices.size() > 0) {
                    Map<String, Object> choice = (Map<String, Object>) choices.get(0);
                    Map<String, Object> msg = (Map<String, Object>) choice.get("message");
                    String content = String.valueOf(msg.get("content"));
                    responseEntity.setSuccess(true);
                    responseEntity.setContent(content);

                    Map<String, Object> usage = (Map<String, Object>) map.get("usage");
                    int prompt_tokens = (int) usage.get("prompt_tokens");
                    int completion_tokens = (int) usage.get("completion_tokens");
                    responseEntity.setInputTokens(prompt_tokens);
                    responseEntity.setOutputTokens(completion_tokens);
                    return;
                }
            }
            responseEntity.setSuccess(false);
            responseEntity.setErrorMessage(message);
        } catch (Exception e) {
            log.error("parse json failed: {}, message: {}", e.getMessage(), message);
            responseEntity.setSuccess(false);
            responseEntity.setErrorMessage(e.getMessage());
        }
    }


    private void handleWithTextPlain(InputStreamReader reader,
                                     OutputStream outputStream, AiResponseEntity responseEntity) throws IOException {
        int length;
        char[] buffer = new char[1024];
        StringBuilder builder = new StringBuilder();
        while ((length = reader.read(buffer, 0, buffer.length)) != -1) {
            String resp = new String(buffer, 0, length);
            log.info("resp: {}", resp);
            builder.append(resp);
            outputStream.write(resp.getBytes(StandardCharsets.UTF_8));
            outputStream.flush();
        }
        // 与handleWithJson不同在于下面的处理不一样，text plain我们认为是没问题的，是成功的
        responseEntity.setSuccess(true);
        responseEntity.setContent(builder.toString());
        responseEntity.setTokenCount(-1);
    }


    private void handleWithStream(InputStreamReader reader, OutputStream outputStream,
                                  AiResponseEntity responseEntity) {
        StringBuilder builder = new StringBuilder();
        String line = null;
        int tokenCount = 0;
        try {
            BufferedReader bufferedReader = new BufferedReader(reader);
            while ((line = bufferedReader.readLine()) != null) {
                log.debug("resp: {}", line);
                if (line.length() > PREFIX_LENGTH && !"data: [DONE]".equals(line)) {
                    builder.append(getContent(line.substring(PREFIX_LENGTH)));
                    // todo 验证此方式计算token是否正确
                    ++tokenCount;
                }
                outputStream.write("event: event\n".getBytes());
                outputStream.write(line.getBytes(StandardCharsets.UTF_8));
                outputStream.write("\n".getBytes());
                outputStream.flush();
            }
        } catch (IOException e) {
            log.error("io exception: {}, line = {}", e.getMessage(), line, e);
            responseEntity.setErrorMessage(e.getMessage() + (line != null ? ("for line: " + line) : ""));
        }
        if (tokenCount > 0) {
            responseEntity.setSuccess(true);
        }
        responseEntity.setContent(builder.toString());
        responseEntity.setTokenCount(tokenCount);
    }

    @SuppressWarnings("unchecked")
    private String getContent(String line) throws JsonProcessingException {
        Map<String, Object> map = objectMapper.readValue(line, new TypeReference<Map<String, Object>>() {
        });
        List<Map<String, Object>> choices = (List<Map<String, Object>>) map.get("choices");
        Map<String, Object> choice = choices.get(0);
        Map<String, Object> delta = (Map<String, Object>) choice.get("delta");
        boolean hasContent = delta.containsKey("content");
        if (hasContent) {
            return (String) delta.get("content");
        }
        return "";
    }

    private String getRequestBody(ChatRequestEntity chatRequestEntity, AiModelEntity modelEntity) {

//        String model = modelEntity.getName();
        List<Map<String, Object>> messages = chatRequestEntity.getMessages().stream()
                .map(msg -> {
                    Map<String, Object> map = new HashMap<>();
                    map.put("role", msg.getRole());
                    map.put("content", msg.getContent());
                    return map;
                }).collect(Collectors.toList());

        // 构建请求参数
        HashMap<Object, Object> params = new HashMap<>();
        params.put("stream", "true".equals(chatRequestEntity.getStream()));
        if(StringUtil.isNotEmpty(chatRequestEntity.getModel())){
            params.put("model", chatRequestEntity.getModel());
        }
        params.put("messages", messages);
        if(ObjectUtil.isNotEmpty(chatRequestEntity.getVariables())){
            params.put("variables", chatRequestEntity.getVariables());
        }
        params.put("presence_penalty", chatRequestEntity.getPresencePenalty());
//        params.put("temperature", chatRequestEntity.getTemperature());

        return JSONUtil.toJsonStr(params);
    }

    private String getBaseUrl(AiModelEntity modelEntity) {
        int platformId = modelEntity.getPlatformId();
        String baseUrl = modelEntity.getPlatformEntity().getBaseUrl();
        if (baseUrl == null || baseUrl.length() == 0) {
            if (platformId == AiPlatformEntity.ID_OPEN_AI) {
                baseUrl = AiPlatformEntity.BASE_URL_OPEN_AI;
            } else {
                log.error("the platform unsupported: platformId = {}", platformId);
                throw new IllegalArgumentException("the platform unsupported");
            }
        }
        String path = modelEntity.getPath();
        if (path == null || path.length() == 0) {
            path = AiModelEntity.PATH_CHAT_GPT;
            //if (platformId == AiPlatformEntity.ID_OPEN_AI) {
            //    path = AiModelEntity.PATH_CHAT_GPT;
            //} else {
            //    log.error("the platform unsupported: platformId = {}", platformId);
            //    throw new IllegalArgumentException("the platform unsupported");
            //}
        }
        if (!path.startsWith("/")) {
            path = "/" + path;
        }
        return baseUrl + path;

    }


    // reference: https://jtokkit.knuddels.de/docs/getting-started/recipes/chatml
    // https://github.com/openai/openai-cookbook/blob/main/examples/How_to_count_tokens_with_tiktoken.ipynb
    private int countMessageTokens(String model, List<ChatRequestEntity.Message> messages) {
        Encoding encoding = registry.getEncodingForModel(model).orElse(null);
        if (encoding == null) {
            return -1; //throw new IllegalArgumentException("Unsupported model for get encoding: " + model);
        }
        int tokensPerMessage;
        int tokensPerName;
        String modelName = model.toLowerCase();
        if (Objects.equals(modelName, "gpt-3.5-turbo-0301")) {
            tokensPerMessage = 4; // every message follows <|start|>{role/name}\n{content}<|end|>\n
            tokensPerName = -1; // if there's a name, the role is omitted
        } else if (modelName.startsWith("gpt-4") || modelName.startsWith("gpt-3.5")) {
            tokensPerMessage = 3;
            tokensPerName = 1;
        } else {
            return -1; //throw new IllegalArgumentException("Unsupported model for count tokens: " + model);
        }

        int sum = 0;
        for (ChatRequestEntity.Message message : messages) {
            sum += tokensPerMessage;
            sum += encoding.countTokens(StringUtils.toStringTrim(message.getContent()));
            sum += encoding.countTokens(message.getRole());
            if (message.isHasName()) {
                sum += encoding.countTokens(message.getName());
                sum += tokensPerName;
            }
        }

        //sum += 3; // every reply is primed with <|start|>assistant<|message|>

        return sum;
    }

    private static final EncodingRegistry registry = Encodings.newDefaultEncodingRegistry();

    private static int getTokens(String text) {
        if (StrUtil.isBlank(text)) {
            return 0;
        } else {
            Encoding enc = registry.getEncoding(EncodingType.CL100K_BASE);
            return enc.encode(text).size();
        }
    }
}
